import time
import requests
import json
from tqdm import tqdm
import threading

import os

# 定义一个类 APIModel，用于通过 API 与模型交互
class APIModel:

    # 初始化方法，设置模型名称、API 密钥和 API URL
    def __init__(self, model, api_key, api_url) -> None:
        self.__api_key = api_key  # 存储 API 密钥（私有属性）
        self.__api_url = api_url  # 存储 API URL（私有属性）
        self.model = model  # 模型名称
        
    # 发送请求的内部方法
    def __req(self, text, temperature, max_try=10):
        url = f"{self.__api_url}/chat/completions"  # OpenAI API 请求的 URL
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.__api_key}',  # 使用 API 密钥进行身份验证
        }
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": text}],
            "temperature": temperature,
        }
        
        for _ in range(max_try):
            try:
                response = requests.post(url, headers=headers, json=payload)
                response.raise_for_status()  # 检查请求是否成功
                return response.json()['choices'][0]['message']['content']  # 返回响应内容
            except requests.exceptions.RequestException as e:
                print(f"请求失败: {e}")
                time.sleep(2)  # 重试前的短暂延迟
        return None  # 如果所有尝试都失败，返回 None
    
    # 与模型进行单次交互的接口方法
    def chat(self, text, temperature=1):
        response = self.__req(text, temperature=temperature, max_try=5)  # 调用内部请求方法获取响应
        return response  # 返回响应内容

    # 用于多线程请求时的内部方法
    def __chat(self, text, temperature, res_l, idx):
        response = self.__req(text, temperature=temperature)  # 获取响应
        res_l[idx] = response  # 将响应存储在指定索引位置
        return response
        
    # 批量发送消息与模型交互的方法，支持多线程
    def batch_chat(self, text_batch, temperature=0):
        max_threads = 20 # 限制同时使用的最大线程数
        res_l = ['No response'] * len(text_batch)  # 初始化响应列表，默认值为 'No response'
        thread_l = []  # 用于存储线程的列表
        # 遍历消息批次并创建线程
        for i, text in zip(range(len(text_batch)), text_batch):
            thread = threading.Thread(target=self.__chat, args=(text, temperature, res_l, i))  # 创建线程
            thread_l.append(thread)  # 将线程添加到列表中
            thread.start()  # 启动线程
            # 控制最大线程数量，防止超出限制
            while len(thread_l) >= max_threads: 
                for t in thread_l:
                    if not t.is_alive():  # 检查线程是否已经结束
                        thread_l.remove(t)  # 如果线程结束，则从列表中移除
                time.sleep(2)  # 短暂延迟以避免忙等待

        # 等待所有线程结束
        # for thread in tqdm(thread_l):
        for thread in thread_l:
            thread.join()  # 等待线程完成
        return res_l  # 返回所有响应  

if __name__ == "__main__":
    # 初始化 APIModel 实例
    model = APIModel(
        model="glm-4-air", 
        api_key='99fbf79da7e241929139d3e9fcc0e3d2.M8PE9i2FJiVd3RBy',
        api_url='https://open.bigmodel.cn/api/paas/v4/'
        )
    
    # model = APIModel(
    #     model="ep-20250227122716-gzj5s", 
    #     api_key='7b448f0e-e6cb-421e-94d0-e232e2540441',
    #     api_url='https://ark.cn-beijing.volces.com/api/v3'
    #     )
    
    # 单次交互示例
    response = model.chat("你好，你是谁？")
    print(f"单次交互的响应: {response}")
    
    # 批量交互示例
    text_batch = ["请介绍一下你自己。", "今天的热点新闻是什么？", "你有什么爱好吗？"]
    responses = model.batch_chat(text_batch)
    for i, res in enumerate(responses):
        print(f"批量交互的第 {i+1} 个响应: {res}")